﻿using System;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Http;
using MySqlX.XDevAPI;
using Newtonsoft.Json;
using Renci.SshNet;



public class SshWebSocketMiddleware
{
    private readonly RequestDelegate _next;

    public SshWebSocketMiddleware(RequestDelegate next)
    {
        _next = next;
    }

    public async Task InvokeAsync(HttpContext context)
    {
        if (context.WebSockets.IsWebSocketRequest)
        {
            using var webSocket = await context.WebSockets.AcceptWebSocketAsync();
            await HandleWebSocketAsync(context, webSocket);
        }
        else
        {
            await _next(context);
        }
    }

    private async Task HandleWebSocketAsync(HttpContext context, WebSocket webSocket)
    {
        var buffer = new byte[1024 * 4];
        var webSocketReceiveResult = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
        var configJson = Encoding.UTF8.GetString(buffer, 0, webSocketReceiveResult.Count);
        using var client = new SshClient("119.45.240.39", "root", "Zt04985807");
        client.Connect();
        try
        {

            

            var shellStream = client.CreateShellStream("xterm", 80, 24, 800, 600, 1024);

            var sshOutputTask = Task.Run(async () =>
            {
                var outputBuffer = new byte[1024];
                while (true)
                {
                    var outputLength = shellStream.Read(outputBuffer, 0, outputBuffer.Length);
                    if (outputLength > 0)
                    {
                        var output = Encoding.UTF8.GetString(outputBuffer, 0, outputLength);
                        var outputBytes = Encoding.UTF8.GetBytes(output);
                        await webSocket.SendAsync(new ArraySegment<byte>(outputBytes, 0, outputBytes.Length), WebSocketMessageType.Text, true, CancellationToken.None);
                    }
                    await Task.Delay(TimeSpan.FromMilliseconds(100));
                }
            });

            while (true)
            {
                webSocketReceiveResult = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);

                if (webSocketReceiveResult.MessageType == WebSocketMessageType.Text)
                {
                    var command = Encoding.UTF8.GetString(buffer, 0, webSocketReceiveResult.Count);
                    shellStream.WriteLine(command);
                }
                else if (webSocketReceiveResult.MessageType == WebSocketMessageType.Close)
                {
                    await webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "", CancellationToken.None);
                    break;
                }
                await Task.Delay(TimeSpan.FromMilliseconds(100));
            }

            await sshOutputTask;

            client.Disconnect();
        }
        catch (Exception ex)
        {
            //_logger.LogError(ex, "An error occurred while handling the SSH WebSocket connection.");
            Serilog.Log.Information($"WebSocket ERROR: {ex.Message}");
            // 发送错误消息到前端，如果WebSocket仍然开启
            if (webSocket.State == WebSocketState.Open)
            {
                var errorMsg = Encoding.UTF8.GetBytes("An internal error occurred.");
                await webSocket.SendAsync(new ArraySegment<byte>(errorMsg), WebSocketMessageType.Text, true, CancellationToken.None);
            }
        }
        finally
        {
            if (client?.IsConnected == true)
            {
                client.Disconnect();
            }
            client?.Dispose();

            if (webSocket?.State == WebSocketState.Open)
            {
                await webSocket.CloseAsync(WebSocketCloseStatus.InternalServerError, "Connection closed by the server.", CancellationToken.None);
            }
        }
    }
}
